/*
 * Copyright (c) <2025> <shihaorui> and <2501060546@qq.com>. All rights reserved.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Contributor: <shihaorui> (<2501060546@qq.com>)
 */
 #include "segment_min_coo_tiling.h"
 #include "register/op_def_registry.h"
 #include "tiling/platform/platform_ascendc.h"
 
 
 namespace optiling {
 const uint32_t BLOCK_SIZE = 32;
 const int32_t NUM = 4;
 static ge::graphStatus TilingFunc(gert::TilingContext* context)
 {
 
   SegmentMinCooTilingData tiling;
   uint32_t sizeofdatatype;
   uint32_t totalLengthAligned;
   auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
   auto socVersion = ascendcPlatform.GetSocVersion();
   uint64_t ubLength;
   ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubLength);
   auto aivCanUseNum = ascendcPlatform.GetCoreNum();
   auto dt = context->GetInputDesc(0)->GetDataType();
   if(dt == ge::DT_INT8){
     sizeofdatatype = 1;
   }else if(dt ==ge::DT_FLOAT16 || dt == ge::DT_BF16){
     sizeofdatatype = 2;
   }else{
     sizeofdatatype = 4;
   }
   const auto& srcShape = context->GetInputShape(0)->GetStorageShape();
   const auto& idxShape = context->GetInputShape(1)->GetStorageShape();
   uint32_t srcDimNum = srcShape.GetDimNum();
   uint32_t idxDimNum = idxShape.GetDimNum();
   uint32_t srcTotalSize = srcShape.GetShapeSize();
   uint32_t indexPhysicalNum = idxShape.GetShapeSize();
   uint32_t E_1 = 1;
   uint32_t logicalM = 1;
   uint32_t idxMPhysical = 1;
   if (idxDimNum >= 1) {
     for (uint32_t d = 0; d + 1 < idxDimNum; ++d) {
       uint32_t s = srcShape.GetDim(d);
       uint32_t i = idxShape.GetDim(d);
       if (!(i == s || i == 1)) {
         return ge::GRAPH_FAILED;
       }
       E_1 *= s;
     }
     logicalM = srcShape.GetDim(idxDimNum - 1);
     idxMPhysical = idxShape.GetDim(idxDimNum - 1);
     if (!(idxMPhysical == logicalM || idxMPhysical == 1)) {
       return ge::GRAPH_FAILED;
     }
   } else {
     E_1 = 1;
     logicalM = 1;
     idxMPhysical = 1;
   }
   tiling.set_E_1(E_1);
   tiling.set_logicalM(logicalM);
   tiling.set_indexPhysicalNum(indexPhysicalNum);
   uint32_t logicalIndexNum = E_1 * logicalM;
   tiling.set_logicalIndexNum(logicalIndexNum);
   if (logicalIndexNum == 0 || srcTotalSize % logicalIndexNum != 0) {
     return ge::GRAPH_FAILED;
   }
   uint32_t K = srcTotalSize / logicalIndexNum;
   tiling.set_K(K);
 
   auto attrs = context->GetAttrs();
   const int32_t* nSegmentsPtr = attrs->GetAttrPointer<int32_t>(0);
   if (nSegmentsPtr == nullptr || *nSegmentsPtr <= 0) {
     return ge::GRAPH_FAILED;
   }
   uint32_t nSegments = static_cast<uint32_t>(*nSegmentsPtr);
   tiling.set_nSegments(nSegments);
 
   tiling.set_indexNum(indexPhysicalNum);
 
   uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
 
   uint32_t ubPartBlockNum = ubLength / BLOCK_SIZE / 2 / NUM;
 
   ubPartBlockNum = ubPartBlockNum <= 8 ? ubPartBlockNum : ubPartBlockNum / 8 * 8;
 
   uint32_t ubPartDataNum = ubPartBlockNum * ALIGN_NUM;
 
   uint32_t aivNum = (aivCanUseNum < logicalIndexNum / ubPartDataNum) ? aivCanUseNum : (logicalIndexNum / ubPartDataNum);
   aivNum = aivNum >= 1 ? aivNum : 1;
 
   uint32_t coreDataNum = 1;
   if(aivNum != 0 && ALIGN_NUM != 0){
 // Align to 256B
     coreDataNum = (logicalIndexNum / aivNum) / (ALIGN_NUM * 8) * (ALIGN_NUM * 8);
   }
   uint32_t coreTailDataNum = logicalIndexNum - aivNum * coreDataNum;
 
   tiling.set_srcLength(srcTotalSize);
   tiling.set_ubPartDataNum(ubPartDataNum);
   tiling.set_coreDataNum(coreDataNum);
   tiling.set_coreTailDataNum(coreTailDataNum);
   tiling.set_ALIGN_NUM(ALIGN_NUM);
   tiling.set_aivNum(aivNum);
   uint32_t strideM = 1;
   if (idxMPhysical == 1 && logicalM > 1) {
     strideM = 0;
   }
   uint32_t idxPreProduct = 1;
   for (uint32_t d = 0; d + 1 < idxDimNum; ++d) {
     idxPreProduct *= idxShape.GetDim(d);
   }
   uint32_t strideE1 = (idxPreProduct == E_1 ? idxMPhysical : 0);
   tiling.set_strideM(strideM);
   tiling.set_strideE1(strideE1);
   
   tiling.set_idxDimNum(idxDimNum);
   
   uint32_t logicalShapeArr[8] = {0};
   uint32_t idxPhysicalShapeArr[8] = {0};
   uint32_t idxStridesArr[8] = {0};
   
   if (idxDimNum > 0 && idxDimNum <= 8) {
     for (uint32_t d = 0; d < idxDimNum; ++d) {
       logicalShapeArr[d] = srcShape.GetDim(d);
       idxPhysicalShapeArr[d] = idxShape.GetDim(d);
     }
     
     for (uint32_t d = 0; d < idxDimNum; ++d) {
       uint32_t physicalDim = idxShape.GetDim(d);
       uint32_t logicalDim = srcShape.GetDim(d);
       
       if (physicalDim == 1 && logicalDim > 1) {
         idxStridesArr[d] = 0;
       } else {
         uint32_t stride = 1;
         for (uint32_t nextD = d + 1; nextD < idxDimNum; ++nextD) {
           stride *= idxShape.GetDim(nextD);
         }
         idxStridesArr[d] = stride;
       }
     }
   }
   
   tiling.set_logicalShape(logicalShapeArr);
   tiling.set_idxPhysicalShape(idxPhysicalShapeArr);
   tiling.set_idxStrides(idxStridesArr);
   
   context->SetBlockDim(aivNum);
   uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
   size_t *currentWorkspace = context->GetWorkspaceSizes(1);
   size_t usrSize = static_cast<size_t>(E_1) * static_cast<size_t>(nSegments) * sizeof(uint8_t);
   currentWorkspace[0] = usrSize + sysWorkspaceSize;
   tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
   context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
 
   return ge::GRAPH_SUCCESS;
 }
 }
 
 namespace ge {
 static ge::graphStatus InferShape(gert::InferShapeContext* context)
 {
     auto attrs = context->GetAttrs();
     if (attrs == nullptr) {
         return GRAPH_FAILED;
     }
     const int32_t* nSegmentsPtr = attrs->GetAttrPointer<int32_t>(0);
     if (nSegmentsPtr == nullptr || *nSegmentsPtr <= 0) {
         return GRAPH_FAILED;
     }
     int32_t nSegments = *nSegmentsPtr;
     const gert::Shape* x1_shape = context->GetInputShape(0);
     const gert::Shape* index_shape = context->GetInputShape(1);
     
     gert::Shape* y_shape = context->GetOutputShape(0);
     *y_shape = *x1_shape;
     int32_t index_last_dim = index_shape->GetDimNum() - 1;
     y_shape->SetDim(index_last_dim, nSegments);
     
     gert::Shape* arg_out_shape = context->GetOutputShape(1);
     *arg_out_shape = *y_shape;
     
     return GRAPH_SUCCESS;
 }
 static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
 {
     const auto inputDataType = context->GetInputDataType(0);
     context->SetOutputDataType(0, inputDataType);
     context->SetOutputDataType(1, ge::DT_INT32);
     return ge::GRAPH_SUCCESS;
 }
 }
 
 
 namespace ops {
 class SegmentMinCoo : public OpDef {
 public:
     explicit SegmentMinCoo(const char* name) : OpDef(name)
     {
         this->Input("src")
             .ParamType(REQUIRED)
             .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
         this->Input("index")
             .ParamType(REQUIRED)
             .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
         this->Output("out")
             .ParamType(REQUIRED)
             .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32})
             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
         this->Output("arg_out")
             .ParamType(REQUIRED)
             .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
             .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
             .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
         this->Attr("nSegments")
             .AttrType(REQUIRED)
             .Int(0);
         this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
 
         this->AICore()
             .SetTiling(optiling::TilingFunc);
         this->AICore().AddConfig("ascend910b");
 
     }
 };
 
 OP_ADD(SegmentMinCoo);
 }